import os
import logging
import pickle
import numpy as np
try:
    import cv2
except ImportError:
    cv2 = None

# 尝试导入paths，如果失败则使用os.walk替代
try:
    import paths
    HAS_PATHS = True
except ImportError:
    HAS_PATHS = False
    import glob

from tqdm import tqdm

# 定义函数get_size,获取文件大小,返回文件大小的兆字节的浮点数表示
def get_size(file):
    """
    获取指定文件的大小（以MB为单位）
    
    参数:
    file (str): 文件的路径
    
    返回:
    float: 文件大小（MB）
    """
    # 获取文件大小（字节），然后转换为MB
    return os.path.getsize(file) / (1024 * 1024)

# 如果没有paths模块，创建替代函数
def list_images(folder):
    """使用glob列出文件夹中的所有图像文件"""
    if not os.path.exists(folder):
        logging.warning(f"文件夹 {folder} 不存在")
        return []
    
    image_extensions = ['*.jpg', '*.jpeg', '*.png', '*.bmp', '*.gif']
    image_paths = []
    for extension in image_extensions:
        image_paths.extend(glob.glob(os.path.join(folder, extension)))
        image_paths.extend(glob.glob(os.path.join(folder, '**', extension), recursive=True))
    return image_paths

# 定义函数createXY,用于从图像创建特征(X)和标签(y)
def createXY(train_folder, dest_folder, method='vgg', batch_size=64):
    x_file_path = os.path.join(dest_folder, "X.pkl")  # 设置X文件的路径,用于保存特征数据
    y_file_path = os.path.join(dest_folder, "y.pkl")  # 设置y文件的路径,用于保存标签数据

    # 如果 X 和 y 已经存在，则直接读取，不再重新构建
    if os.path.exists(x_file_path) and os.path.exists(y_file_path):
        logging.info("X和y已经存在，直接读取")  # 提示用户X和y文件已经存在
        logging.info(f"X文件大小:{get_size(x_file_path):.2f}MB")  # 打印X文件的大小
        logging.info(f"y文件大小:{get_size(y_file_path):.2f}MB")  # 打印y文件的大小
        
        # 用 pickle 读取X和y文件
        with open(x_file_path, 'rb') as f:
            X = pickle.load(f)
        with open(y_file_path, 'rb') as f:
            y = pickle.load(f)
        return X, y  # 直接返回已加载的数据

    if cv2 is None:
        raise ImportError("需要安装opencv-python库: pip install opencv-python")

    logging.info("读取所有图像，生成X和y")  # 提示用户开始读取图像并生成X和y
    
    # 根据是否有paths模块选择不同的图像路径获取方法
    if HAS_PATHS:
        image_paths = list(paths.list_images(train_folder))  # 获取训练文件夹中所有图像的路径
    else:
        image_paths = list_images(train_folder)  # 使用替代方法获取图像路径
    
    logging.info(f"找到 {len(image_paths)} 张图像")

    if len(image_paths) == 0:
        raise ValueError(f"在文件夹 {train_folder} 中未找到任何图像文件")

    X = []  # 初始化X列表,用于存储特征
    y = []  # 初始化y列表,用于存储标签
    
    # 计算需要的批次数
    num_batches = len(image_paths) // batch_size + (1 if len(image_paths) % batch_size else 0)

    # 使用进度条对批次进行循环处理
    for idx in tqdm(range(num_batches), desc="读取图像"):
        batch_images = []  # 初始化存储批次图像的列表
        batch_labels = []  # 初始化存储批次标签的列表
        
        start = idx * batch_size  # 计算批次开始的索引
        end = min((idx + 1) * batch_size, len(image_paths))  # 计算批次结束的索引

        # 对于每个批次中的图像
        for i in range(start, end):
            image_path = image_paths[i]  # 获取图像路径
            # 不再区分vgg和flat方法，统一使用简单图像处理
            img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)  # 以灰度模式读取图像
            if img is not None:  # 确保图像读取成功
                img = cv2.resize(img, (32, 32))  # 调整图像大小到32x32
                batch_images.append(img)  # 将图像数组添加到批次图像列表
                
                label = os.path.basename(image_path).split('.')[0]  # 从文件名中解析出字符串标签
                label = 1 if label == 'dog' else 0  # 如果标签是'dog'则为1,否则为0
                batch_labels.extend([label])  # 将标签添加到批次标签列表

        if len(batch_images) > 0:
            batch_images = np.array(batch_images)  # 将批次图像列表转换为numpy数组
            batch_pixels = batch_images.reshape(batch_images.shape[0], -1)  # 将 batch_images 展平

            X.extend(batch_pixels)  # 将处理后的图像特征添加到X列表
            y.extend(batch_labels)  # 将标签添加到y列表

    logging.info(f"X.shape: {np.shape(X)}")  # 打印X的形状
    logging.info(f"y.shape: {np.shape(y)}")  # 打印y的形状
    
    if len(X) == 0:
        raise ValueError("未能从图像中提取任何特征，请检查图像文件是否有效")

    # 将X和y分别序列化到文件
    with open(x_file_path, 'wb') as f:
        pickle.dump(X, f)
    with open(y_file_path, 'wb') as f:
        pickle.dump(y, f)

    return X, y  # 返回构建的X和y